"""
可视化模块
功能：生成损伤热力图和波形对比图
"""
import matplotlib.pyplot as plt
import seaborn as sns

def plot_waveform_comparison(original_wave, processed_wave, save_path):
    """
    绘制原始波形与处理后波形对比
    参数：
    - original_wave: 原始三轴波形 [3, 1000]
    - processed_wave: 预处理后波形
    - save_path: 图片保存路径
    """
    plt.figure(figsize=(12, 6))
    for i, axis in enumerate(['X', 'Y', 'Z']):
        plt.subplot(3, 1, i+1)
        plt.plot(original_wave[i], label='原始信号', alpha=0.7)
        plt.plot(processed_wave[i], label='预处理后', linestyle='--')
        plt.title(f"{axis}轴波形对比")
        plt.legend()
    plt.tight_layout()
    plt.savefig(save_path)
    plt.close()

def plot_damage_heatmap(floor_probs, save_path):
    """
    生成楼层损伤概率热力图
    参数：
    - floor_probs: 各楼层概率列表（如[0.2, 0.5, 0.3]）
    - save_path: 图片保存路径
    """
    plt.figure(figsize=(8, 6))
    sns.heatmap([floor_probs], annot=True, fmt=".2f", cmap="YlGnBu",
                xticklabels=[f"Floor {i+1}" for i in range(len(floor_probs))])
    plt.title("各楼层损伤概率分布")
    plt.savefig(save_path)
    plt.close()